# Load MNIST data
import torch
import torchvision
import torchvision.datasets
import torchvision.transforms as transforms
from torch.utils.data import Subset

data_tf = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
)

train_dataset = torchvision.datasets.MNIST(
    root="./MNIST", train=True, transform=data_tf, download=True
)

test_dataset = torchvision.datasets.MNIST(
    root="./MNIST", train=False, transform=data_tf, download=True
)

print(train_dataset.targets.size(), train_dataset.data.size())

# Filter indices of the train and test dataset for label 8
train_indices = torch.where(train_dataset.targets == 8)[0]
test_indices = torch.where(test_dataset.targets == 8)[0]

# Create subsets for label 8
train_dataset_label_8 = Subset(train_dataset, train_indices)  # type: ignore
test_dataset_label_8 = Subset(test_dataset, test_indices)  # type: ignore

# Verify the size of the new datasets
print(len(train_dataset_label_8), len(test_dataset_label_8))


dataloaders = {
    "train": torch.utils.data.DataLoader(  # type: ignore
        train_dataset, batch_size=128, shuffle=True, num_workers=3
    ),
    "test": torch.utils.data.DataLoader(  # type: ignore
        test_dataset, batch_size=128, shuffle=False, num_workers=3
    ),
    "train_label_8": torch.utils.data.DataLoader(  # type: ignore
        train_dataset_label_8, batch_size=128, shuffle=True, num_workers=3
    ),
    "test_label_8": torch.utils.data.DataLoader(  # type: ignore
        test_dataset_label_8, batch_size=128, shuffle=False, num_workers=3
    ),
}


def get_dataloader(dataset, batch_size):
    if dataset == "train":
        return torch.utils.data.DataLoader(  # type: ignore
            train_dataset, batch_size=batch_size, shuffle=True, num_workers=3
        )
    elif dataset == "test":
        return torch.utils.data.DataLoader(  # type: ignore
            test_dataset, batch_size=batch_size, shuffle=False, num_workers=3
        )


# import matplotlib.pyplot as plt
# plt.imshow(train_dataset.data[0].numpy(), cmap='gray')
# plt.title('%i' % train_dataset.targets[0])
# plt.savefig('MNIST_example.png')
